# -*- coding:utf-8 -*-
# CREATED BY: jiangbohuai
# CREATED ON:  4:05 PM
# LAST MODIFIED ON:
# AIM: 一个抽象类
import abc
import os

import torch
from torch import nn
from typing import Optional, Tuple
import matplotlib.pyplot as plt
import numpy as np

from nn.loss.loss_abc import LossABC


class TorchModelABC(nn.Module):
    def __init__(self,
                 output_path: str,
                 loss: LossABC):
        super(TorchModelABC, self).__init__()
        self.output_path = output_path
        if not os.path.isdir(output_path):
            os.mkdir(output_path)

        self.loss = loss
        # -- global parameters -- #
        self.hist_train_loss = []
        self.average_train = []

        self.hist_valid_loss = []
        self.average_valid = []

        self.best_loss = np.inf

    def plot_performance(self, path: str):
        plt.figure()
        self.average_train.append(float(np.mean(self.hist_train_loss)))
        plt.plot(self.average_train, label='train loss')
        if self.hist_valid_loss:
            self.average_valid.append(float(np.mean(self.hist_valid_loss)))
            plt.plot(self.average_valid, label='valid loss')
            best_loss = np.min(self.average_valid)
        else:
            best_loss = np.min(self.average_train)

        plt.savefig(os.path.join(self.output_path, path))
        plt.legend()
        plt.close()

        self.hist_train_loss = []
        self.hist_valid_loss = []

        if best_loss < self.best_loss:
            self.best_loss = best_loss
            self.save('best_model')

    def get_loss(self,
                 actual,
                 expect,
                 act_valid: Optional = None,
                 exp_valid: Optional = None) -> Tuple[torch.Tensor, str]:
        loss_train, verbose_train = self.loss(actual, expect)
        self.hist_train_loss.append(loss_train.item())
        if act_valid is not None:
            loss_valid, verbose_valid = self.loss(act_valid, exp_valid)
            self.hist_valid_loss.append(loss_valid.item())
            verbose = f'trian {verbose_train} valid {verbose_valid}'
        else:
            verbose = verbose_train
        return loss_train, verbose

    def save(self, path: str):
        torch.save(self.state_dict(), os.path.join(self.output_path, path))

    def load(self, path: str):
        self.load_state_dict(torch.load(path))

    def load_best_model(self):
        self.load( os.path.join(self.output_path, 'best_model'))

    def predict(self, *args, **kwargs):
        return self.forward(*args, **kwargs)

    @abc.abstractmethod
    def forward(self, *args, **kwargs):
        pass

    @abc.abstractmethod
    def train_on_batch(self,
                       plot_path: str,
                       verbose: bool,
                       *args,
                       **kwargs):
        pass
